import sys
import os
import re
from llama_cpp import Llama
from PyQt5.QtWidgets import (QApplication, QWidget, QVBoxLayout, QHBoxLayout, QLabel,
                             QTextEdit, QPushButton, QRadioButton, QMessageBox,
                             QButtonGroup, QComboBox, QCheckBox, QSplitter)
from PyQt5.QtCore import Qt, QTimer, QThread, pyqtSignal
from PyQt5.QtGui import QFont

# by DZgas 
# v3 03.03.2026 

LANGUAGES = [
    "Russian", "English", "German", "Spanish", "French", "Italian", "Korean", "Dutch",
    "Portuguese (Portugal)", "Portuguese (Brazilian)", "Spanish (Latin America)",
    "Chinese (Simplified)", "Chinese (Traditional)", "Czech", "Ukrainian", "Hindi",
    "Icelandic", "Japanese", "Polish", "Swedish", "Hungarian", "Romanian",
    "Danish", "Norwegian (Nynorsk)", "Norwegian (Bokmål)", "Finnish"
]
COMBO_PLACEHOLDER = "Другой язык..."

def create_prompt_content(text, target_lang_name):
    return (
        f"Translate the following text to {target_lang_name}:\n"
        f"{text}\n"
        f"{target_lang_name}: "
    )

def initialize_model(model_path):
    print("\n---- debug ----")
    try:
        llm = Llama(
            model_path=model_path,
            n_gpu_layers=0,
            n_ctx=4096,          
            verbose=False,
            cache=True,       
            swa_full=True,     
        )

        max_tokens = llm.n_ctx()
        
        check_text = "hey guys"
        prompt_content = create_prompt_content(check_text, "Russian")
        messages = [{"role": "user", "content": prompt_content}]
        
        output = llm.create_chat_completion(
            messages=messages,
            max_tokens=15,
            temperature=0.0
        )
        result = output['choices'][0]['message']['content']
        if not result.strip():
             raise ValueError("Model check returned an empty result.")

        return llm, max_tokens

    except Exception as e:
        import traceback
        traceback.print_exc()
        return None, None

class TranslatorWorker(QThread):
    finished = pyqtSignal(str, bool) 

    def __init__(self, pipe, messages, generation_args):
        super().__init__()
        self.pipe = pipe
        self.messages = messages
        self.generation_args = generation_args
        self.input_text_snapshot = ""
        self.input_tokens_len = 0

    def run(self):
        try:
            outputs = self.pipe.create_chat_completion(
                messages=self.messages,
                **self.generation_args
            )
            raw_result = outputs['choices'][0]['message']['content']
            final_result = raw_result.strip()
            self.finished.emit(final_result, True)
        except Exception as e:
            self.finished.emit(str(e), False)

class TranslationApp(QWidget):
    
    def __init__(self, llm_instance, max_tokens):
        super().__init__()
        self.pipe = llm_instance
        self.max_tokens = max_tokens
        
        self.last_translated_text = ""
        self.last_translated_token_count = 0
        
        self.debounce_timer = QTimer()
        self.debounce_timer.setSingleShot(True)
        self.debounce_timer.setInterval(1500) 
        self.debounce_timer.timeout.connect(self.start_translation)
        
        self.current_worker = None
        self.init_ui()

    def init_ui(self):
        self.setWindowTitle("Нейросетевой оффлайн переводчик")
        self.setGeometry(300, 300, 800, 500)
        self.setMinimumSize(100, 124) 
        
        font = QFont("Arial", 12)
        
        lang_settings_layout = QHBoxLayout()
        lang_settings_layout.addWidget(QLabel("Перевести на:"))
        
        self.ru_radio = QRadioButton("Русский")
        self.en_radio = QRadioButton("Английский")
        self.lang_button_group = QButtonGroup()
        self.lang_button_group.addButton(self.ru_radio)
        self.lang_button_group.addButton(self.en_radio)
        lang_settings_layout.addWidget(self.ru_radio)
        lang_settings_layout.addWidget(self.en_radio)
        
        self.target_lang_combo = QComboBox()
        self.target_lang_combo.addItem(COMBO_PLACEHOLDER)
        self.target_lang_combo.addItems(sorted([lang for lang in LANGUAGES if lang not in ["Russian", "English"]]))
        lang_settings_layout.addWidget(self.target_lang_combo)
        
        self.auto_detect_check = QCheckBox("Автосмена (RU/EN)")
        self.auto_detect_check.setChecked(True)
        lang_settings_layout.addWidget(self.auto_detect_check)
        
        lang_settings_layout.addStretch()
        self.ru_radio.setChecked(True)
        
        temp_layout = QHBoxLayout()
        temp_layout.addWidget(QLabel("Разнообразие:"))
        self.temp_button_group = QButtonGroup()
        self.temp0_radio = QRadioButton("0.0")
        self.temp0_radio.setChecked(True)
        self.temp05_radio = QRadioButton("0.5")
        self.temp1_radio = QRadioButton("1.0")
        self.temp_button_group.addButton(self.temp0_radio)
        self.temp_button_group.addButton(self.temp05_radio)
        self.temp_button_group.addButton(self.temp1_radio)
        temp_layout.addWidget(self.temp0_radio)
        temp_layout.addWidget(self.temp05_radio)
        temp_layout.addWidget(self.temp1_radio)
        temp_layout.addStretch()
        
        self.input_text = QTextEdit()
        self.input_text.setFont(font)
        self.input_text.setPlaceholderText("Ввод...")
        
        self.output_text = QTextEdit()
        self.output_text.setFont(font)
        self.output_text.setReadOnly(True)
        
        self.token_count_label = QLabel(f"Tokens: 0/{self.max_tokens}")
        self.status_label = QLabel("Готов")
        
        self.copy_button = QPushButton("📋")
        self.copy_button.setEnabled(False)
        self.copy_button.setMaximumWidth(25)
        
        self.swap_button = QPushButton("⇅")
        self.swap_button.setMaximumWidth(25)
        
        main_layout = QVBoxLayout()
        
        top_bar = QHBoxLayout()
        top_bar.addLayout(lang_settings_layout, 2)
        top_bar.addLayout(temp_layout, 1)
        main_layout.addLayout(top_bar)
        
        left_layout = QVBoxLayout()
        left_layout.addWidget(self.input_text)
        left_bottom = QHBoxLayout()
        left_bottom.addWidget(self.token_count_label)
        left_bottom.addWidget(self.swap_button)
        
        self.layout_check = QCheckBox("〓")
        self.layout_check.toggled.connect(self.change_layout_orientation)
        left_bottom.addWidget(self.layout_check)
        
        self.dark_mode_check = QCheckBox("🌙") # Галочка с луной для тёмной темы
        self.dark_mode_check.toggled.connect(self.toggle_dark_theme)
        left_bottom.addWidget(self.dark_mode_check)
        
        left_bottom.addStretch()
        left_layout.addLayout(left_bottom)
        
        left_widget = QWidget()
        left_widget.setLayout(left_layout)
        
        right_layout = QVBoxLayout()
        right_layout.addWidget(self.output_text)
        right_bottom = QHBoxLayout()
        right_bottom.addWidget(self.status_label)
        right_bottom.addStretch()
        right_bottom.addWidget(self.copy_button)
        right_layout.addLayout(right_bottom)
        
        right_widget = QWidget()
        right_widget.setLayout(right_layout)
        
        # ИЗМЕНЕНИЕ 1: Сохраняем сплиттер в self, чтобы менять его позже
        self.main_splitter = QSplitter(Qt.Horizontal)
        self.main_splitter.addWidget(left_widget)
        self.main_splitter.addWidget(right_widget)
        self.main_splitter.setStretchFactor(0, 1)
        self.main_splitter.setStretchFactor(1, 1)
        self.main_splitter.setCollapsible(0, False)
        self.main_splitter.setCollapsible(1, False)
        
        main_layout.addWidget(self.main_splitter)
        

        
        self.setLayout(main_layout)
        
        # СВЯЗИ СОБЫТИЙ
        self.input_text.textChanged.connect(self.on_text_change)
        self.swap_button.clicked.connect(self.swap_languages_and_text)
        self.copy_button.clicked.connect(self.copy_output_to_clipboard)
        
        self.ru_radio.clicked.connect(self.on_language_changed_manual)
        self.en_radio.clicked.connect(self.on_language_changed_manual)
        self.target_lang_combo.currentIndexChanged.connect(self.on_combo_changed)
        
        self.ru_radio.toggled.connect(self.on_language_changed_manual)
        self.en_radio.toggled.connect(self.on_language_changed_manual)
        
        self.temp_button_group.buttonClicked.connect(self.on_temperature_changed)
        self.auto_detect_check.clicked.connect(self.on_auto_detect_toggled)
        
        # ИЗМЕНЕНИЕ 3: Связь для смены вида
        self.layout_check.toggled.connect(self.change_layout_orientation)
        
        self.input_text.setFocus(Qt.OtherFocusReason)
        
        
    def toggle_dark_theme(self, checked):
        if checked:
            # Простой стиль тёмной темы
            self.setStyleSheet("""
                QWidget {
                    background-color: #2b2b2b;
                    color: #ffffff;
                }
                QTextEdit {
                    background-color: #1e1e1e;
                    color: #FFFFFF;
                    border: 1px solid #3d3d3d;
                }
                QPushButton {
                    background-color: #3d3d3d;
                    color: #ffffff;
                    border: 1px solid #555;
                    padding: 3px;
                }
                QPushButton:hover {
                    background-color: #4d4d4d;
                }
                QComboBox {
                    background-color: #3d3d3d;
                    color: #ffffff;
                    border: 1px solid #555;
                }
                QSplitter::handle {
                    background-color: #2b2b2b;
                }
            """)
        else:
            # Возврат к стандартной теме
            self.setStyleSheet("")
            
            
    def detect_language_and_switch(self, text):
        if not self.auto_detect_check.isChecked() or not text.strip():
            return

        cyrillic_count = len(re.findall(r'[а-яА-Я]', text))
        latin_count = len(re.findall(r'[a-zA-Z]', text))

        if cyrillic_count == 0 and latin_count == 0:
            return

        need_signal_block = False
        
        if cyrillic_count > latin_count:
            # Ввод Русский -> Цель Английский
            if not self.en_radio.isChecked():
                self.en_radio.setChecked(True)
                self.target_lang_combo.setCurrentIndex(0)
                need_signal_block = True
                
        elif latin_count > cyrillic_count:
            # Ввод Английский -> Цель Русский
            if not self.ru_radio.isChecked():
                self.ru_radio.setChecked(True)
                self.target_lang_combo.setCurrentIndex(0)
                need_signal_block = True
                
        if need_signal_block:
             pass 

    def on_language_changed_manual(self):
        # Реагируем только если кнопка была ВКЛЮЧЕНА (чтобы не срабатывало дважды при переключении)
        if not self.sender().isChecked():
            return

        # Сбрасываем выпадающий список
        self.target_lang_combo.blockSignals(True)
        self.target_lang_combo.setCurrentIndex(0)
        self.target_lang_combo.blockSignals(False)
        
        # СБРОС ПАМЯТИ: Заставляем переводить заново
        self.last_translated_text = "" 
        self.start_translation()

    def on_temperature_changed(self):
        # СБРОС ПАМЯТИ: Температура изменилась -> результат будет другим
        self.last_translated_text = ""
        self.start_translation()

    def on_auto_detect_toggled(self):
        # Если включили галочку - сразу пробуем определить язык текущего текста
        if self.auto_detect_check.isChecked():
            text = self.input_text.toPlainText()
            self.detect_language_and_switch(text)
        
        # СБРОС ПАМЯТИ и перезапуск
        self.last_translated_text = ""
        self.start_translation()

    def on_combo_changed(self):
        # Если выбран реальный язык из списка
        if self.target_lang_combo.currentText() != COMBO_PLACEHOLDER:
            # Отключаем автосмену и радио-кнопки
            self.auto_detect_check.setChecked(False)
            self.lang_button_group.setExclusive(False)
            self.ru_radio.setChecked(False)
            self.en_radio.setChecked(False)
            self.lang_button_group.setExclusive(True)
            
            # ВАЖНО: Сбрасываем "память"
            self.last_translated_text = ""
            
            self.start_translation()

    def on_text_change(self):
        input_content = self.input_text.toPlainText()
        
        self.detect_language_and_switch(input_content)
        
        target_lang = self.get_target_language()
        is_valid_state = bool(input_content.strip() and target_lang)
        
        if not is_valid_state:
            self.token_count_label.setText(f"Tokens: 0/{self.max_tokens}")
            self.status_label.setText("Ожидание...")
            return
            
        try:
            prompt_content = create_prompt_content(input_content, target_lang)
            input_ids = self.pipe.tokenize(prompt_content.encode("utf-8"))
            current_token_count = len(input_ids)
        except Exception:
            current_token_count = 0
            
        if current_token_count > self.max_tokens:
            self.token_count_label.setStyleSheet("color: red;")
            self.status_label.setText("Лимит!")
        else:
            self.token_count_label.setStyleSheet("color: black;")
            self.token_count_label.setText(f"Tokens: {current_token_count}/{self.max_tokens}")
            
            self.debounce_timer.start()

            token_diff = abs(current_token_count - self.last_translated_token_count)
            
            if token_diff >= 20:
                if not self.current_worker or not self.current_worker.isRunning():
                    self.status_label.setText("Порог 20...")
                    self.start_translation()

    def start_translation(self):
        if self.current_worker and self.current_worker.isRunning():
            return

        input_content = self.input_text.toPlainText()
        
        if len(input_content) < 2: 
            return
        
        if not input_content.strip() or input_content == self.last_translated_text:
            return

        target_lang = self.get_target_language()
        if not target_lang:
            self.status_label.setText("Нет языка")
            return
            
        prompt_content = create_prompt_content(input_content, target_lang)
        try:
            input_ids = self.pipe.tokenize(prompt_content.encode("utf-8"))
            input_tokens_count = len(input_ids)
        except: return

        if input_tokens_count > self.max_tokens:
            self.status_label.setText("Много текста")
            return

        max_new = self.max_tokens - input_tokens_count - 10
        if max_new <= 0:
            return

        temp_text = self.temp_button_group.checkedButton().text()
        temperature = float(re.findall(r"(\d+\.\d+)", temp_text)[0])
        
        generation_args = {
            "max_tokens": max(input_tokens_count * 2, 100), 
            "temperature": temperature,
            "stop": ["User:", "Translate the following text"]
        }
        
        self.status_label.setText("Перевод...")
        
        messages = [{"role": "user", "content": prompt_content}]
        self.current_worker = TranslatorWorker(self.pipe, messages, generation_args)
        
        self.current_worker.input_tokens_len = input_tokens_count 
        self.current_worker.input_text_snapshot = input_content 
        
        self.current_worker.finished.connect(self.on_translation_finished)
        self.current_worker.start()

    def on_translation_finished(self, result, success):
        worker = self.sender()
        if not worker: return # Защита от случайного вызова
        
        if success:
            self.output_text.setText(result)
            self.status_label.setText("Готово")
            self.copy_button.setEnabled(bool(result.strip()))
            
            self.last_translated_text = worker.input_text_snapshot
            self.last_translated_token_count = worker.input_tokens_len
            
            current_text = self.input_text.toPlainText()
            if current_text != self.last_translated_text:
                self.debounce_timer.start()
        else:
            self.output_text.setText(f"Ошибка: {result}")
            self.status_label.setText("Ошибка")
        
        self.current_worker = None

    def get_target_language(self):
        if self.ru_radio.isChecked(): return "Russian"
        if self.en_radio.isChecked(): return "English"
        combo_text = self.target_lang_combo.currentText()
        if combo_text != COMBO_PLACEHOLDER: return combo_text
        return None

    def copy_output_to_clipboard(self):
        QApplication.clipboard().setText(self.output_text.toPlainText())

    def swap_languages_and_text(self):
        self.auto_detect_check.setChecked(False) 
        
        current_target_is_ru = self.ru_radio.isChecked()
        current_target_is_en = self.en_radio.isChecked()
        
        if current_target_is_ru:
            self.en_radio.setChecked(True)
        elif current_target_is_en:
            self.ru_radio.setChecked(True)

        input_content = self.input_text.toPlainText()
        output_content = self.output_text.toPlainText()
        
        self.input_text.blockSignals(True)
        self.input_text.setText(output_content)
        self.input_text.blockSignals(False)
        self.output_text.setText(input_content)
        
        self.last_translated_text = "" 
        self.start_translation()

    def change_layout_orientation(self, checked):
        # Если галочка нажата (checked) -> Вертикально, иначе -> Горизонтально
        if checked:
            self.main_splitter.setOrientation(Qt.Vertical)
        else:
            self.main_splitter.setOrientation(Qt.Horizontal)

    def closeEvent(self, event):
        # Этот метод вызывается автоматически при закрытии окна (на крестик)
        if self.current_worker and self.current_worker.isRunning():
            self.current_worker.quit()
            self.current_worker.wait(1000) # Ждем 1 секунду завершения потока
        event.accept()


if __name__ == "__main__":
    app = QApplication(sys.argv)
    
    # 1. Автоматически ищем GGUF файл в текущей папке
    try:
        # Находим все файлы с расширением .gguf (без учета регистра)
        gguf_files = sorted([f for f in os.listdir('.') if f.lower().endswith('.gguf')])
        if not gguf_files:
            raise FileNotFoundError("В папке со скриптом не найдено ни одного GGUF файла.")
        
        # Берем первый файл по алфавиту
        model_path = gguf_files[0]

    except FileNotFoundError as e:
        # Если модель не найдена, показываем ошибку и выходим
        error_msg = QMessageBox()
        error_msg.setIcon(QMessageBox.Critical)
        error_msg.setText("Модель не найдена")
        error_msg.setInformativeText(str(e))
        error_msg.exec_()
        sys.exit(1)

    # 2. Выводим информацию о загружаемой модели
    print(f"Loading model: {model_path}")
    
    # Загружаем модель (эта функция теперь будет выводить разделитель)
    llm, max_tokens = initialize_model(model_path)
    
    if llm:
        window = TranslationApp(llm, max_tokens)
        window.show()
        sys.exit(app.exec_())
    else:
        error_msg = QMessageBox()
        error_msg.setIcon(QMessageBox.Critical)
        error_msg.setText("Критическая ошибка")
        error_msg.setInformativeText(f"Не удалось загрузить модель: {model_path}")
        error_msg.exec_()
        sys.exit(1)